#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import xml.etree.ElementTree as ET
from dataclasses import dataclass
from typing import List, Optional
from datetime import datetime

@dataclass
class Danmaku:
    """弹幕信息"""
    time: float      # 弹幕出现时间（秒）
    type: int       # 弹幕类型
    color: int      # 弹幕颜色
    text: str       # 弹幕内容
    user_id: str    # 用户ID
    timestamp: int  # 发送时间戳

class DanmakuUtil:
    @staticmethod
    def parse_xml(xml_content: str) -> List[Danmaku]:
        """解析XML格式的弹幕文件"""
        root = ET.fromstring(xml_content)
        danmaku_list = []
        
        for d in root.findall('.//d'):
            attrs = d.attrib['p'].split(',')
            danmaku = Danmaku(
                time=float(attrs[0]),
                type=int(attrs[1]),
                color=int(attrs[3]),
                text=d.text,
                user_id=attrs[6],
                timestamp=int(attrs[4])
            )
            danmaku_list.append(danmaku)
        
        return danmaku_list
    
    @staticmethod
    def convert_to_ass(danmaku_list: List[Danmaku], output_path: str) -> None:
        """将弹幕转换为ASS格式"""
        with open(output_path, 'w', encoding='utf-8') as f:
            # 写入ASS头部
            f.write("""[Script Info]
ScriptType: v4.00+
PlayResX: 1920
PlayResY: 1080
Timer: 100.0000

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,微软雅黑,54,&H00FFFFFF,&H00FFFFFF,&H00000000,&H00000000,0,0,0,0,100,100,0,0,1,2,0,2,20,20,20,1

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
""")
            
            # 写入弹幕
            for d in danmaku_list:
                start_time = datetime.fromtimestamp(d.timestamp).strftime('%H:%M:%S.%f')[:-3]
                end_time = datetime.fromtimestamp(d.timestamp + 5).strftime('%H:%M:%S.%f')[:-3]
                color = f"&H{d.color:08X}"
                f.write(f"Dialogue: 0,{start_time},{end_time},Default,,0,0,0,,{{\\c{color}}}{d.text}\n")
    
    @staticmethod
    def convert_to_srt(danmaku_list: List[Danmaku], output_path: str) -> None:
        """将弹幕转换为SRT格式"""
        with open(output_path, 'w', encoding='utf-8') as f:
            for i, d in enumerate(danmaku_list, 1):
                start_time = datetime.fromtimestamp(d.timestamp).strftime('%H:%M:%S,%f')[:-3]
                end_time = datetime.fromtimestamp(d.timestamp + 5).strftime('%H:%M:%S,%f')[:-3]
                f.write(f"{i}\n{start_time} --> {end_time}\n{d.text}\n\n") 